package io.grpc.benchmarks.qps.service;

import com.google.protobuf.ByteString;
import io.grpc.Status;
import io.grpc.stub.StreamObserver;
import io.grpc.testing.Payload;
import io.grpc.testing.PayloadType;
import io.grpc.testing.SimpleRequest;
import io.grpc.testing.SimpleResponse;
import io.grpc.testing.TestServiceGrpc;

public class TestServiceImpl implements TestServiceGrpc.TestService {

    @Override
    public void unaryCall(SimpleRequest request,
            StreamObserver<SimpleResponse> responseObserver) {
        SimpleResponse response = buildSimpleResponse(request);
        responseObserver.onNext(response);
        responseObserver.onCompleted();
    }

    @Override
    public StreamObserver<SimpleRequest> streamingCall(
            final StreamObserver<SimpleResponse> responseObserver) {
        return new StreamObserver<SimpleRequest>() {
            @Override
            public void onNext(SimpleRequest request) {
                SimpleResponse response = buildSimpleResponse(request);
                responseObserver.onNext(response);
            }

            @Override
            public void onError(Throwable t) {
                System.out.println("Encountered an error in streamingCall");
                t.printStackTrace();
            }

            @Override
            public void onCompleted() {
                responseObserver.onCompleted();
            }
        };
    }

    private static SimpleResponse buildSimpleResponse(SimpleRequest request) {
        if (request.getResponseSize() > 0) {
            if (!PayloadType.COMPRESSABLE.equals(request.getResponseType())) {
                throw Status.INTERNAL
                    .augmentDescription("Error creating payload.")
                    .asRuntimeException();
            }

            ByteString body = ByteString
                .copyFrom(new byte[request.getResponseSize()]);
            PayloadType type = request.getResponseType();

            Payload payload = Payload.newBuilder().setType(type).setBody(body)
                .build();
            return SimpleResponse.newBuilder().setPayload(payload).build();
        }
        return SimpleResponse.getDefaultInstance();
    }
}